import os
import PIL

from torchvision import transforms
from torch.utils.data import Dataset, DataLoader


class ImageDataset(Dataset):
    def __init__(self, filelist, transform=None):
        super(ImageDataset, self).__init__()
        self.filelist = filelist
        self.transform = transform

    def __len__(self):
        return len(self.filelist)

    def __getitem__(self, index):
        image = PIL.Image.open(self.filelist[index]).convert('RGB')
        name = self.filelist[index].split('/')[-1]
        image = self.transform(image) \
            if self.transform is not None \
            else transforms.ToTensor(image)
        return image, name
    

# class FeatureDataset(Dataset):
#     def __init__(self, data):
#         super(FeatureDataset, self).__init__()


def load_data():
    # Small test
    # inpdir = '/home/xincz/Documents/yy_search_test/卡通/todo'
    # tardir = '/home/xincz/Documents/yy_search_test/卡通/ori'

    # Medium test
    # inpdir = '/home/xincz/Documents/yy_search_test/small_queries'
    # tardir = '/home/xincz/Documents/yy_search_test/small_pool'

    # Large test
    inpdir = '/home/xincz/Documents/yy_search_test/query_100'
    tardir = '/home/xincz/Documents/yy_search_test/pool_1w'
    inp_files = os.listdir(inpdir)
    tar_files = os.listdir(tardir)

    def fullpath(filelist, folderpath):
        return ['/'.join([folderpath, filename]) for filename in filelist]

    inp_files = fullpath(inp_files, inpdir)
    tar_files = fullpath(tar_files, tardir)

    input_count = len(inp_files)
    target_count = len(tar_files)

    print("input count: ", input_count)
    print("target count: ", target_count)
    
    inp_trans = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
    ])

    tar_trans = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
    ])

    inp_ds = ImageDataset(inp_files, inp_trans)
    inp_dl = DataLoader(inp_ds, shuffle=False, num_workers=0)

    tar_ds = ImageDataset(tar_files, tar_trans)
    tar_dl = DataLoader(tar_ds, shuffle=False, batch_size=16, num_workers=4)
    
    return {
        "input_count": input_count,
        "target_count": target_count,
        "inp_ds": inp_ds,
        "tar_dl": tar_dl,
    }
